import numpy as np
import torch


# 索引 、切片
# 2:   2到最后个
# :2    第一个到第2个
# ::2   每隔1个
print("索引获取数据----------->")
print(torch.randn(2, 3, 2, 2).shape)
print(list(torch.randn(2, 3, 2, 2).shape))
print(torch.randn(2, 3, 2, 2)[0])
print(torch.randn(2, 3, 2, 2)[0].shape)

data = torch.randn(2, 3, 6, 6)

print(data[1][2])
# 等价于
print(data[1, 2])
print(data[1, 2, 1, 0])
# ：2 正序从第一个到2个
print(data[:2, :, :, :].shape)
# 后面的 ：可以不写
print(data[:2].shape)

# -1索引倒序到最后一个
print(data[:1, -1:].shape)
# 1: 索引第一个到最后一个
print(data[:1, 1:].shape)
# 起始：结尾：步长  或者 ：：2
print(data[:, :, 0:6:2, 0:6:2].shape)
print(data[:, :, ::2, ::2])

# 索引获取
print(data.index_select(0, torch.tensor([0, 1])))

# .代表任意多
d = torch.randn(5, 3, 28, 28)
print(d[...].shape)
print(d[1, ...].shape)
print(d[:, 1, ...].shape)
print(d[..., 2:].shape)

# mask掩码
print("mask掩码---------------->")
x = torch.randn(3, 4)
print(x)
# 大于0.5
mask = x.ge(0.5)
# 小于
# x.lt()

print(mask)
print(torch.masked_select(x, mask=mask))

y = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
print(torch.take(y, torch.tensor([0, 1, 2])))
